import os
import sys
def benign_file_process(benign_file):
    with open(benign_file, 'r', encoding='utf-8') as f:
        benign_code = f.readlines()

    include_stat = []
    body_stat = []
    main_stat = []
    main_flag = 0
    for line in benign_code:
        line_splited = line.split(' ')
        if main_flag == 1:
            if len(line) > 0 and line[0] == '}':
                main_flag = 0
            else:
                main_stat.append(line)
        elif line_splited[0] == '#include':
            include_stat.append(line)
        elif len(line_splited) < 2 or len(line_splited[1]) < 4 or line_splited[1][0:4] != 'main':
            body_stat.append(line)
        else:
            main_flag = 1
    return [include_stat, body_stat, main_stat]

def juliet_file_process(juliet_file):
    with open(juliet_file, 'r', encoding='utf-8') as f:
        juliet_code = f.readlines()

    func_stat = []
    main_stat = []
    success_flag = False
    for line in juliet_code:
        if '#ifdef INCLUDEMAIN' not in line:
            func_stat.append(line)
        else:
            success_flag = True
            for i in range(juliet_code.index(line), len(juliet_code)):
                # if 'printLine("Calling good()...");' in juliet_code[i] or 'printLine("Calling bad()...");' in juliet_code[i]:
                #     pass
                # elif 'printLine("Finished good()");' in juliet_code[i] or 'printLine("Finished bad()");' in juliet_code[i]:
                #     main_stat.append('/* replace_point */')
                if 'printLine("Calling good()...");' in juliet_code[i] or 'printLine("Calling bad()...");' in juliet_code[i]:
                    pass
                elif 'printLine("Finished bad()");' in juliet_code[i] or 'printLine("Finished good()");' in juliet_code[i]:
                    main_stat.append('/* replace_point */')
                else:
                    main_stat.append(juliet_code[i])
            break
    if success_flag == True:
        return [func_stat, main_stat]
    else:
        return 'pass'

def simple_merge(benign_dir, juliet_dir):
    benign_files = [f for f in os.listdir(benign_dir) if f.endswith('.c')]
    juliet_files = [f for f in os.listdir(juliet_dir) if (f.endswith('.c') or f.endswith('.cpp'))]

    if not benign_files or not juliet_files:
        print("Error: No file found.")
        return
    
    benign_files_num = len(benign_files)
    juliet_files_num = len(juliet_files)

    for i in range(juliet_files_num):
        benign_file = benign_files[i % benign_files_num]
        juliet_file = juliet_files[i]

        benign_file_path = os.path.join(benign_dir, benign_file)
        juliet_file_path = os.path.join(juliet_dir, juliet_file)

        result = []
        benign_code = benign_file_process(benign_file_path)
        juliet_code = juliet_file_process(juliet_file_path)

        if juliet_code == 'pass':
            print(f'No replace point in {juliet_file}')
            # os.remove(juliet_file_path)
            # print(f'{juliet_file} removed')
            continue
        
        result = benign_code[0]
        result.append('\n')
        result.extend(juliet_code[0])
        result.append('\n')
        result.extend(benign_code[1])
        result.append('\n')
        for line in juliet_code[1]:
            if '/* replace_point */' not in line:
                result.append(line)
            else:
                result.extend(benign_code[2])
                result.append('\n')
        
        with open(juliet_file_path, 'w', encoding='utf-8') as file:
            file.writelines(result)

        print(f'Updated {juliet_file} with content from {benign_file}')

# benign_dir = '/home/kokomi/static_project/benign'
# juliet_dir = '/home/kokomi/static_project/workspace/s02_1'

def main():

    simple_merge(benign_dir, juliet_dir)


if __name__ == '__main__':

    if len(sys.argv) != 3:
        print("用法: python simple_merge.py <benign_path> <juliet_path> 或者 不输入参数默认使用脚本路径")

        if len(sys.argv) == 1:
            benign_dir = '/home/kokomi/static_project/workspace/test1'
            juliet_dir = '/home/kokomi/static_project/workspace/s1000'
            main()
    else:
        benign_dir = sys.argv[1]
        juliet_dir = sys.argv[2]
        main()



